import argparse
import logging
import os
from typing import Tuple

import cv2
import numpy as np
import torch

from models.cmnext_conf_3 import CMNeXtWithConf
from models.modal_extract import ModalitiesExtractor
import torchvision.transforms.functional as TF
import torch
from torch import nn, Tensor
from models.layers import trunc_normal_
from torch.nn import functional as F


def load_dualpath_model(model, model_file, backbone):
    extra_pretrained = model_file if 'MHSA' in backbone else None
    if isinstance(extra_pretrained, str):
        raw_state_dict_ext = torch.load(extra_pretrained, map_location=torch.device('cpu'))
        if 'state_dict' in raw_state_dict_ext.keys():
            raw_state_dict_ext = raw_state_dict_ext['state_dict']
    if isinstance(model_file, str):
        raw_state_dict = torch.load(model_file, map_location=torch.device('cpu'))
        if 'model' in raw_state_dict.keys():
            raw_state_dict = raw_state_dict['model']
    else:
        raw_state_dict = model_file

    state_dict = {}
    for k, v in raw_state_dict.items():
        if k.find('patch_embed') >= 0:
            state_dict[k] = v
        elif k.find('block') >= 0:
            state_dict[k] = v
        elif k.find('norm') >= 0:
            state_dict[k] = v

    if isinstance(extra_pretrained, str):
        for k, v in raw_state_dict_ext.items():
            if k.find('patch_embed1.proj') >= 0:
                state_dict[k.replace('patch_embed1.proj', 'extra_downsample_layers.0.proj.module')] = v
            if k.find('patch_embed2.proj') >= 0:
                state_dict[k.replace('patch_embed2.proj', 'extra_downsample_layers.1.proj.module')] = v
            if k.find('patch_embed3.proj') >= 0:
                state_dict[k.replace('patch_embed3.proj', 'extra_downsample_layers.2.proj.module')] = v
            if k.find('patch_embed4.proj') >= 0:
                state_dict[k.replace('patch_embed4.proj', 'extra_downsample_layers.3.proj.module')] = v

            if k.find('patch_embed1.norm') >= 0:
                for i in range(model.num_modals):
                    state_dict[k.replace('patch_embed1.norm', 'extra_downsample_layers.0.norm.ln_{}'.format(i))] = v
            if k.find('patch_embed2.norm') >= 0:
                for i in range(model.num_modals):
                    state_dict[k.replace('patch_embed2.norm', 'extra_downsample_layers.1.norm.ln_{}'.format(i))] = v
            if k.find('patch_embed3.norm') >= 0:
                for i in range(model.num_modals):
                    state_dict[k.replace('patch_embed3.norm', 'extra_downsample_layers.2.norm.ln_{}'.format(i))] = v
            if k.find('patch_embed4.norm') >= 0:
                for i in range(model.num_modals):
                    state_dict[k.replace('patch_embed4.norm', 'extra_downsample_layers.3.norm.ln_{}'.format(i))] = v
            elif k.find('block') >= 0:
                state_dict[k.replace('block', 'extra_block')] = v
            elif k.find('norm') >= 0:
                state_dict[k.replace('norm', 'extra_norm')] = v

    msg = model.load_state_dict(state_dict, strict=False)
    del state_dict


def weighted_statistics_pooling(x, log_w=None):
    b = x.shape[0]
    c = x.shape[1]
    x = x.view(b, c, -1)

    if log_w is None:
        log_w = torch.zeros((b, 1, x.shape[-1]), device=x.device)
    else:
        assert log_w.shape[0] == b
        assert log_w.shape[1] == 1
        log_w = log_w.view(b, 1, -1)

        assert log_w.shape[-1] == x.shape[-1]

    log_w = F.log_softmax(log_w, dim=-1)
    x_min = -torch.logsumexp(log_w - x, dim=-1)
    x_max = torch.logsumexp(log_w + x, dim=-1)

    w = torch.exp(log_w)
    x_avg = torch.sum(w * x, dim=-1)
    x_msq = torch.sum(w * x * x, dim=-1)

    x = torch.cat((x_min, x_max, x_avg, x_msq), dim=1)

    return x


#####
def test(image_path,ckpt_path):
    os.environ['CUDA_VISIBLE_DEVICES'] = '1'
    device = 'cuda:1'
    if device != 'cpu':
        # cudnn setting
        import torch.backends.cudnn as cudnn

        cudnn.benchmark = False
        cudnn.deterministic = True
        cudnn.enabled = True
    # modal_extractor = ModalitiesExtractor(list(('noiseprint', 'bayar', 'srm')), 'pretrained/noiseprint/np++.pth')
    # modal_extractor = ModalitiesExtractor(list(('noiseprint')), 'pretrained/noiseprint/np++.pth')

    from configs.cmnext_init_cfg import _C as config , update_config
    parser = argparse.ArgumentParser(description='')
    # parser.add_argument('-gpu', '--gpu', type=int, default=0, help='device, use -1 for cpu')
    parser.add_argument('-log', '--log', type=str, default='INFO', help='logging level')
    parser.add_argument('-train_bayar', '--train_bayar', default=False, action='store_true', help='finetune bayar conv')
    parser.add_argument('-exp', '--exp', type=str, default='experiments/ec_example.yaml', help='Yaml experiment file')
    parser.add_argument('opts', help="other options", default=None, nargs=argparse.REMAINDER)

    args = parser.parse_args()

    config = update_config(config, args.exp)
    model = CMNeXtWithConf(config.MODEL)

    modal_extractor = ModalitiesExtractor(config.MODEL.MODALS[1:], config.MODEL.NP_WEIGHTS)
    ckpt = torch.load(ckpt_path, map_location="cpu")  # .pth 文件路径
    modal_extractor.to(device)
    model = model.to(device)
    model.load_state_dict(ckpt['state_dict'], strict=True)
    modal_extractor.load_state_dict(ckpt['extractor_state_dict'])


    modal_extractor.set_val()
    model.set_val()

    image = cv2.cvtColor(cv2.imread(image_path), cv2.COLOR_BGR2RGB)
    h, w, c = image.shape
    image_transforms_final = A.Compose([
        ToTensorV2()
    ])

    if h > 2048 or w > 2048:
        res = A.LongestMaxSize(max_size=1024)(image=image, mask=None)
        image = res['image']

    image = image_transforms_final(image=image)['image']
    image = image / 256.0
    image = image.to(device)
    image = image.unsqueeze(0)
    # print(image.shape)
    modals = modal_extractor(image)
    images_norm = TF.normalize(image, mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
    inp = [images_norm] + modals
    if config.MODEL.AAM:
        pred, _ = model(inp,0,None)
    # pred = model(inp)
    if config.MODEL.ENM:
        pred = model(inp, 0)

    # map = torch.nn.functional.softmax(pred, dim=1)[:, 1, :, :].squeeze().cpu().numpy()
    map = torch.nn.functional.softmax(pred, dim=1)[:, 1:2, :, :]

    return map, image

def get_contour(label):
    # lbl = label.gt(0.5).float()
    ero = 1 - F.max_pool2d(1 - label, kernel_size=5, stride=1, padding=2)  # erosion
    dil = F.max_pool2d(label, kernel_size=5, stride=1, padding=2)  # dilation

    edge = dil - ero
    return edge
if __name__ == '__main__':
    os.environ['CUDA_VISIBLE_DEVICES'] = '1'
    import albumentations as A
    from albumentations.pytorch import ToTensorV2
    import torch
    import matplotlib.pyplot as plt
    # ckpt_path = "/home/wc/disk1/MMFusion/ckpt/TruFor_btm_de_vit/best_val_loss.pth"
    # ckpt_path = "/home/wc/disk1/MMFusion/ckpt/TruFor_ori/final.pth"
    # ckpt_path = "/home/wc/disk1/MMFusion/ckpt/TruFor_btm_vit_v5/final.pth"  # ours
    # ckpt_path = "/home/wc/disk1/MMFusion/ckpt/TruFor/final.pth"   # trufor
    ckpt_path = "/home/wc/disk1/MMFusion/ckpt/TruFor_AAM_2/final.pth"   # AAM
    # ckpt_path = "/home/wc/disk1/MMFusion/ckpt/TruFor_AAM_aug/final.pth"   # AAM + aug
    # ckpt_path = "/home/wc/disk1/MMFusion/ckpt/TruFor_btm/best_val_loss.pth"   # ENM
    # map = test("/home/wc/disk1/datasets/CASIAV1/Modified_Tp/CM/Sp_S_NNN_A_arc0022_arc0022_0015.jpg", ckpt_path)
    # map = test("/home/wc/disk1/datasets/CASIAV1/Modified_Tp/CM/Sp_S_NNN_A_nat0037_nat0037_0207.jpg", ckpt_path)
    # map = test("/home/wc/disk1/datasets/COVER/tampered/31t.tif", ckpt_path)
    # map = test("/home/wc/disk1/datasets/COVER/tampered/91t.tif", ckpt_path)
    # map = test("/home/wc/disk1/datasets/COVER/tampered/79t.tif", ckpt_path)
    # map = test("/home/wc/disk1/datasets/CASIAV2/image_all/Tamper/Tp/Tp_D_CRN_M_B_nat10165_nat10164_12100.jpg", ckpt_path)
    ##  热力图可视化
    # pred, image = test("/home/wc/disk1/datasets/CASIAV2/image_all/Tamper/Tp/Tp_D_CNN_M_N_cha00026_cha00028_11784.jpg", ckpt_path)
    # pred, image = test("/home/wc/disk1/datasets/CASIAV2/image_all/Tamper/Tp/Tp_D_CNN_S_B_txt00055_txt00047_11328.jpg", ckpt_path)
    # pred, image = test("/home/wc/disk1/datasets/CASIAV2/image_all/Tamper/Tp/Tp_D_CNN_S_O_nat10153_ani00097_12135.jpg", ckpt_path)
    # pred, image = test("/home/wc/disk1/datasets/COLUMBIA/Columbia/4cam_splc/canong3_kodakdcs330_sub_03.tif", ckpt_path)
    # pred, image = test("/home/wc/disk1/datasets/COVER/tampered/16t.tif", ckpt_path)
    # pred, image = test("/home/wc/disk1/datasets/COVER/tampered/43t.tif", ckpt_path)

    # pred, image = test("/home/wc/disk1/datasets/COVER/tampered/28t.tif", ckpt_path)
    # pred, image = test("/home/wc/disk1/datasets/COVER/tampered/30t.tif", ckpt_path)
    # pred, image = test("/home/wc/disk1/datasets/COVER/tampered/31t.tif", ckpt_path)
    # pred, image = test("/home/wc/disk1/datasets/COVER/tampered/34t.tif", ckpt_path)
    # pred, image = test("/home/wc/disk1/datasets/COVER/tampered/35t.tif", ckpt_path)
    # pred, image = test("/home/wc/disk1/datasets/COVER/tampered/42t.tif", ckpt_path)
    pred, image = test("/home/wc/disk1/datasets/COVER/tampered/50t.tif", ckpt_path)
    # pred, image = test("/home/wc/disk1/datasets/COVER/tampered/52t.tif", ckpt_path)
    # pred, image = test("/home/wc/disk1/datasets/COVER/tampered/58t.tif", ckpt_path)
    # pred, image = test("/home/wc/disk1/datasets/COVER/tampered/60t.tif", ckpt_path)
    # pred, image = test("/home/wc/disk1/datasets/COVER/tampered/62t.tif", ckpt_path)

    # 热力图
    # pred, image = test("/home/wc/disk1/datasets/COVER/tampered/65t.tif", ckpt_path)
    # pred, image = test("/home/wc/disk1/datasets/COVER/tampered/67t.tif", ckpt_path)
    # pred, image = test("/home/wc/disk1/datasets/COVER/tampered/70t.tif", ckpt_path)
    # pred, image = test("/home/wc/disk1/datasets/COVER/tampered/77t.tif", ckpt_path)
    # pred, image = test("/home/wc/disk1/datasets/CASIAV1/Modified_Tp/Sp/Sp_D_CNN_A_ani0049_ani0084_0266.jpg", ckpt_path)
    # pred, image = test("/home/wc/disk1/datasets/CASIAV1/Modified_Tp/Sp/Sp_D_CND_A_pla0005_pla0023_0281.jpg", ckpt_path)
    # pred, image = test("/home/wc/disk1/datasets/CASIAV1/Modified_Tp/Sp/Sp_D_CND_A_sec0056_sec0015_0282.jpg", ckpt_path)
    # pred, image = test("/home/wc/disk1/datasets/CASIAV1/Modified_Tp/Sp/Sp_D_CNN_A_ani0053_ani0054_0267.jpg", ckpt_path)
    # pred, image = test("/home/wc/disk1/datasets/CASIAV1/Modified_Tp/Sp/Sp_D_CNN_A_art0024_ani0032_0268.jpg", ckpt_path)
    # pred, image = test("/home/wc/disk1/datasets/CASIAV1/Modified_Tp/Sp/Sp_D_CNN_A_cha0025_pla0067_0269.jpg", ckpt_path)
    # pred, image = test("/home/wc/disk1/datasets/CASIAV1/Modified_Tp/CM/Sp_S_CND_A_pla0016_pla0016_0196.jpg", ckpt_path)
    # pred, image = test("/home/wc/disk1/datasets/CASIAV1/Modified_Tp/CM/Sp_S_CNN_A_pla0052_pla0052_0005.jpg", ckpt_path)
    # pred, image = test("/home/wc/disk1/datasets/CASIAV1/Modified_Tp/CM/Sp_S_CNN_A_pla0084_pla0084_0194.jpg", ckpt_path)
    # pred, image = test("/home/wc/disk1/datasets/CASIAV1/Modified_Tp/CM/Sp_S_CNN_R_sec0097_sec0097_0006.jpg", ckpt_path)
    # pred, image = test("/home/wc/disk1/datasets/CASIAV1/Modified_Tp/Sp/Sp_D_CNN_A_cha0025_pla0067_0269.jpg", ckpt_path)
    # pred, image = test("/home/wc/disk1/datasets/COVER/tampered/92t.tif", ckpt_path)
    # pred, image = test("/home/wc/disk1/datasets/CocoGlide/fake/glide_inpainting_val2017_458702_up.png", ckpt_path)
    # pred, image = test("/home/wc/disk1/datasets/COLUMBIA/Columbia/4cam_splc/canong3_canonxt_sub_30.tif", ckpt_path)
    # pred, image = test("/home/wc/disk1/datasets/COLUMBIA/Columbia/4cam_splc/canong3_kodakdcs330_sub_17.tif", ckpt_path)
    # pred, image = test("/home/wc/disk1/datasets/COLUMBIA/Columbia/4cam_splc/nikond70_canonxt_sub_09.tif", ckpt_path)


    #  xxxxxxx   #
    # pred, image = test("/home/wc/disk1/datasets/CASIAV2/image_all/Tamper/Tp/Tp_D_CRN_M_N_ani10120_sec00098_11632.jpg", ckpt_path)
    # pred, image = test("/home/wc/disk1/datasets/CASIAV2/image_all/Tamper/Tp/Tp_D_CRN_M_N_cha00050_cha00026_11787.jpg", ckpt_path)
    # pred, image = test("/home/wc/disk1/datasets/CASIAV2/image_all/Tamper/Tp/Tp_D_CRN_M_N_nat10156_ani00001_12021.jpg", ckpt_path)
    # pred, image = test("/home/wc/disk1/datasets/CASIAV2/image_all/Tamper/Tp/Tp_D_CRN_S_N_cha10150_cha00040_12221.jpg", ckpt_path)
    # pred, image = test("/home/wc/disk1/datasets/CASIAV2/image_all/Tamper/Tp/Tp_D_CRN_S_N_nat00004_art00025_11426.jpg", ckpt_path) x
    # pred, image = test("/home/wc/disk1/datasets/CASIAV2/image_all/Tamper/Tp/Tp_D_CRN_S_B_pla00071_pla00070_11214.jpg", ckpt_path) x

    # map = test("/home/wc/disk1/datasets/CASIAV1/Modified_Tp/CM/Sp_S_NRD_A_nat0069_nat0069_0258.jpg", ckpt_path) # ×
    # map = test("/home/wc/disk1/datasets/COVER/tampered/27t.tif", ckpt_path)
    # map = test("/home/wc/disk1/datasets/COLUMBIA/Columbia/4cam_splc/canong3_kodakdcs330_sub_25.tif", ckpt_path)  # ×

    pred = pred.cpu().detach().squeeze().numpy()
    ## 可视化边缘
    # edge = get_contour(pred)
    # # 将边缘图从张量转换为 numpy 数组
    # edge_np = edge.squeeze().cpu().detach().numpy()
    #
    # # 可视化边缘图
    # plt.imshow(edge_np, cmap='jet')
    # plt.title('Edge Map')
    # # plt.colorbar()
    # plt.show()

    # ## 可视化热力图
    # heatmap = cv2.applyColorMap(np.uint8(255 * pred), cv2.COLORMAP_PLASMA)  # 生成热力图
    # # heatmap = cv2.applyColorMap(np.uint8(255 * pred), cv2.COLORMAP_COOL)  # 生成热力图
    #
    # image = np.array(image.squeeze().cpu())
    # image = (image * 255).astype(np.uint8)
    # image = np.transpose(image, (1, 2, 0))
    # image = cv2.cvtColor(image, cv2.COLOR_RGB2BGR)
    #
    # # 将热力图叠加到原图上
    # superimposed_img = cv2.addWeighted(image, 0.5, heatmap, 0.5, 0)
    #
    # # 显示结果
    # # plt.figure(figsize=(10, 5))
    # # plt.subplot(1, 2, 1)
    # # plt.title("Original Image")
    # # plt.imshow(cv2.cvtColor(image, cv2.COLOR_BGR2RGB))
    #
    # # plt.subplot(1, 2, 2)
    # plt.figure(figsize=(10, 10))
    # plt.title("Attention Map")
    # plt.imshow(cv2.cvtColor(superimposed_img, cv2.COLOR_BGR2RGB))
    #
    # plt.show()


    # # 可视化图像
    plt.figure(figsize=(10, 10))
    plt.imshow(pred, cmap='gray')  # 使用灰度图显示
    plt.axis('off')  # 关闭坐标轴
    plt.title('Visualizing')
    plt.show()
    # path = "/raid/datasets/ImageForgery/CASIAv2/Tp/Tp_S_NRN_S_N_txt00082_txt00082_11293.jpg"
    # img_path = path.split('_')[-3]
    # img_path = 'Au_' + img_path[0:3] + '_' + img_path[3:] + '.jpg'
    # print(img_path)
    # path = "/raid/datasets/ImageForgery/FantasticReality_v1/dataset/ColorFakeImages/IMG_0000064_IMG_0000063.jpg"
    # path = "/raid/datasets/ImageForgery/CAT-Net/tampCOCO/sp_images/100_000000003517.jpg_000000408480.jpg_RT31.6.jpg"
    # ori = path.split('/')[-1].split('.jpg')[0].split('_')[1]
    # print(ori)
    # ori_path = "/raid/datasets/PoseEstimation/coco/train2017/" + ori + ".jpg"
    # ori_path = "/raid/datasets/ImageForgery/FantasticReality_v1/dataset/ColorRealImages/" + ori[-2] + '_' + ori[-1]
    # ori = path.split('/')[-1].split(".jpg")[0].split('_')[1] + ".jpg"
    # ori_path = "/raid/datasets/PoseEstimation/coco/train2017/" + ori
    # print(ori_path)